"""
股票模块
"""

import read_data
import datetime
import strategy

stock_data = read_data.read_data()  # 股票完整数据
del stock_data["002558.XSHE"]
del stock_data["603110.XSHG"]
del stock_data["002611.XSHE"]
stock_list = list(stock_data.keys())  # 股票代码列表


class Stock(object):
    # 初始化方法，创建一个股票实例
    def __init__(self, code):
        self.code = code
        # 股票自2010-2020的数据
        self.data = stock_data[code]

    # 获取股票代码
    def get_code(self):
        return self.code

    """
    以下的get方法，获取的可能是nan
    """

    # 获取指定日期当天的开盘价
    def get_open(self, date):
        return self.data.loc[date, :].loc["open"]

    # 获取指定日期当天的收盘价，主要用来计算当日收益
    def get_close(self, date):
        return self.data.loc[date, :].loc["close"]

    # 获得前一日的收盘价，计算指标用
    def get_pre_close(self, date):
        return self.data.loc[date, :].loc["pre_close"]

    # 不考虑，手续费,只计算单价，不考虑数量， 第二天收盘后计算
    def cal_total_return(self, day1, day2):
        price1 = self.get_pre_close(day1)
        price2 = self.get_close(day2)
        total_return = (price2 - price1) / price1
        return total_return

    # 计算单只股票的收益率（每日、累计），不考虑手续费,事后计算
    def cal_today_return(self, today):
        price1 = self.get_pre_close(today)  # 昨日收盘
        price2 = self.get_close(today)  # 今日收盘价
        today_return = (price2 - price1) / price1
        return today_return

    def cal_n_day_return(self, today, n):
        price2 = self.get_pre_close(today)
        today_dt = datetime.datetime.strptime(today, "%Y-%m-%d")
        day1_dt = today_dt + datetime.timedelta(days=-n)
        while day1_dt not in strategy.can_trade_date_list_dt:
            day1_dt = day1_dt + datetime.timedelta(days=-1)
        day1 = day1_dt.strftime("%Y-%m-%d")
        price1 = self.get_pre_close(day1)
        return_n = (price2 - price1) / price1
        return return_n


stock_cls = [Stock(code) for code in stock_list]  # 股票实例列表
# print(len(stock_list))
# print(len(stock_cls))
